{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Self-Supervised Learning with Lightly Train\n",
    "\n",
    "[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/opengeos/geoai/blob/main/docs/examples/lightly_train_self_supervised.ipynb)\n",
    "\n",
    "This notebook demonstrates how to use [Lightly Train](https://docs.lightly.ai/train/stable/index.html) for self-supervised pretraining on unlabeled geospatial imagery. You'll learn how to:\n",
    "\n",
    "1. Train a self-supervised model using various methods (DINOv2, DINO, SimCLR)\n",
    "2. Load and use the pretrained model\n",
    "3. Generate embeddings for downstream tasks\n",
    "\n",
    "Self-supervised learning is particularly useful for geospatial applications where labeled data is scarce or expensive to obtain.\n",
    "\n",
    "## Install packages\n",
    "To use the `geoai-py` package with Lightly Train support, ensure it is installed in your environment. Uncomment the command below if needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %pip install geoai-py lightly-train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import geoai\n",
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Download sample data\n",
    "\n",
    "We'll use unlabeled satellite imagery for self-supervised pretraining. For this example, we'll download some sample NAIP imagery."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download sample geospatial imagery for training\n",
    "sample_urls = [\n",
    "    \"https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_rgb_train.tif\",\n",
    "    \"https://huggingface.co/datasets/giswqs/geospatial/resolve/main/cars_7cm.tif\",\n",
    "]\n",
    "\n",
    "# Create directories for our workflow\n",
    "data_dir = \"lightly_train_data\"\n",
    "output_dir = \"lightly_train_output\"\n",
    "embeddings_dir = \"lightly_embeddings\"\n",
    "\n",
    "os.makedirs(data_dir, exist_ok=True)\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "os.makedirs(embeddings_dir, exist_ok=True)\n",
    "\n",
    "# Download and prepare training images\n",
    "image_paths = []\n",
    "for url in sample_urls:\n",
    "    image_path = geoai.download_file(url)\n",
    "    image_paths.append(image_path)\n",
    "    print(f\"Downloaded: {image_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "## Prepare training data\n",
    "\n",
    "For self-supervised learning, we need to extract image patches from our geospatial imagery. These patches will be used as unlabeled training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Extract image patches for training\n",
    "# for i, image_path in enumerate(image_paths):\n",
    "#     patch_output_dir = f\"{data_dir}/patches_{i}\"\n",
    "\n",
    "#     # Extract patches from the raster\n",
    "#     geoai.export_geotiff_tiles(\n",
    "#         in_raster=image_path,\n",
    "#         out_folder=patch_output_dir,\n",
    "#         tile_size=224,  # Common size for vision models\n",
    "#         stride=112,  # 50% overlap\n",
    "#     )\n",
    "\n",
    "#     print(f\"Extracted patches from {image_path} to {patch_output_dir}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Count total training images\n",
    "import glob\n",
    "\n",
    "all_patches = glob.glob(f\"{data_dir}/**/images/*.tif\", recursive=True)\n",
    "print(f\"\\nTotal training patches: {len(all_patches)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Visualize sample training data\n",
    "\n",
    "Let's take a look at some of the extracted patches that will be used for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "\n",
    "# Display a few sample patches\n",
    "sample_patches = all_patches[:6]  # Take first 6 patches\n",
    "\n",
    "fig, axes = plt.subplots(2, 3, figsize=(12, 8))\n",
    "fig.suptitle(\"Sample Training Patches for Self-Supervised Learning\", fontsize=16)\n",
    "\n",
    "for i, patch_path in enumerate(sample_patches):\n",
    "    row, col = i // 3, i % 3\n",
    "    img = Image.open(patch_path)\n",
    "    axes[row, col].imshow(img)\n",
    "    axes[row, col].set_title(f\"Patch {i+1}\")\n",
    "    axes[row, col].axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(f\"Each patch is {img.size} pixels\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## Train self-supervised model\n",
    "\n",
    "Now we'll train a self-supervised model using Lightly Train. We'll use the **SimCLR** method, which works well with CNN models like ResNet.\n",
    "\n",
    "**Method compatibility:**\n",
    "- **SimCLR** or **DINO**: Works with CNN models (ResNet, EfficientNet, etc.)\n",
    "- **DINOv2**: Requires Vision Transformer (ViT) models only\n",
    "\n",
    "**Note**: Training will take some time depending on your hardware. For demonstration purposes, we're using a small number of epochs. In practice, you might want to train for 100+ epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train a self-supervised model using Lightly Train\n",
    "model_path = geoai.lightly_train_model(\n",
    "    data_dir=f\"{data_dir}/patches_0/images\",  # Use patches from first image\n",
    "    output_dir=output_dir,\n",
    "    model=\"torchvision/resnet50\",  # Base architecture\n",
    "    method=\"simclr\",  # Use SimCLR for CNN models like ResNet\n",
    "    epochs=10,  # Small number for demo\n",
    "    batch_size=32,  # Adjust based on your GPU\n",
    "    optim_args={\"lr\": 1e-4},  # Pass learning rate through optim_args\n",
    "    overwrite=True,\n",
    ")\n",
    "\n",
    "print(f\"\\nModel training completed! Pretrained model saved to: {model_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "## Load pretrained model\n",
    "\n",
    "Once training is complete, we can load the pretrained model for use in downstream tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the pretrained model\n",
    "pretrained_model = geoai.load_lightly_pretrained_model(\n",
    "    model_path=model_path, model_architecture=\"torchvision/resnet50\"\n",
    ")\n",
    "\n",
    "print(f\"Loaded pretrained model: {type(pretrained_model)}\")\n",
    "print(f\"Model parameters: {sum(p.numel() for p in pretrained_model.parameters()):,}\")\n",
    "\n",
    "# The model is now ready for fine-tuning on your specific task\n",
    "print(\"\\nModel is ready for fine-tuning on downstream tasks!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "## Generate embeddings\n",
    "\n",
    "We can also use the pretrained model checkpoint to generate embeddings for our images. These embeddings capture rich representations learned through self-supervised training.\n",
    "\n",
    "**Note**: We need to use the checkpoint file (.ckpt) for embedding generation, not the exported model (.pt)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate embeddings for our images\n",
    "# Note: We use the checkpoint file (.ckpt) for embedding generation\n",
    "checkpoint_path = os.path.join(output_dir, \"checkpoints\", \"last.ckpt\")\n",
    "\n",
    "embeddings_path = geoai.lightly_embed_images(\n",
    "    data_dir=f\"{data_dir}/patches_0/images\",\n",
    "    model_path=checkpoint_path,\n",
    "    output_path=f\"{embeddings_dir}/image_embeddings.pt\",\n",
    "    batch_size=32,\n",
    "    overwrite=True,\n",
    ")\n",
    "\n",
    "print(f\"Embeddings saved to: {embeddings_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "## Analyze embeddings\n",
    "\n",
    "Let's load and analyze the generated embeddings to understand what our model has learned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and analyze the embeddings\n",
    "if os.path.exists(embeddings_path):\n",
    "    # Load PyTorch tensor embeddings\n",
    "    embeddings = torch.load(embeddings_path)\n",
    "\n",
    "    # Convert to numpy for analysis if it's a tensor\n",
    "    if isinstance(embeddings, torch.Tensor):\n",
    "        embeddings_np = embeddings.cpu().numpy()[\"embeddings\"]\n",
    "    else:\n",
    "        embeddings_np = embeddings[\"embeddings\"]\n",
    "\n",
    "    print(f\"Embeddings shape: {embeddings_np.shape}\")\n",
    "    print(f\"Embedding dimension: {embeddings_np.shape[1]}\")\n",
    "    print(f\"Number of images embedded: {embeddings_np.shape[0]}\")\n",
    "\n",
    "    # Basic statistics\n",
    "    print(f\"\\nEmbedding statistics:\")\n",
    "    print(f\"Mean: {embeddings_np.mean():.4f}\")\n",
    "    print(f\"Std: {embeddings_np.std():.4f}\")\n",
    "    print(f\"Min: {embeddings_np.min():.4f}\")\n",
    "    print(f\"Max: {embeddings_np.max():.4f}\")\n",
    "else:\n",
    "    print(\n",
    "        \"Embeddings file not found. This might be due to the embedding generation process.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "## Visualize embeddings with t-SNE\n",
    "\n",
    "Let's use t-SNE to visualize the high-dimensional embeddings in 2D space. This helps us understand how well the model groups similar images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize embeddings using t-SNE\n",
    "if os.path.exists(embeddings_path):\n",
    "    try:\n",
    "        from sklearn.manifold import TSNE\n",
    "        import matplotlib.pyplot as plt\n",
    "\n",
    "        # Apply t-SNE for dimensionality reduction\n",
    "        tsne = TSNE(\n",
    "            n_components=2, random_state=42, perplexity=min(30, len(embeddings_np) - 1)\n",
    "        )\n",
    "        embeddings_2d = tsne.fit_transform(embeddings_np)\n",
    "\n",
    "        # Create visualization\n",
    "        plt.figure(figsize=(10, 8))\n",
    "        scatter = plt.scatter(\n",
    "            embeddings_2d[:, 0],\n",
    "            embeddings_2d[:, 1],\n",
    "            c=range(len(embeddings_2d)),\n",
    "            cmap=\"viridis\",\n",
    "            alpha=0.7,\n",
    "        )\n",
    "        plt.colorbar(scatter)\n",
    "        plt.title(\"t-SNE Visualization of Self-Supervised Embeddings\")\n",
    "        plt.xlabel(\"t-SNE Component 1\")\n",
    "        plt.ylabel(\"t-SNE Component 2\")\n",
    "        plt.grid(True, alpha=0.3)\n",
    "        plt.show()\n",
    "\n",
    "        print(\"t-SNE visualization shows how the model groups similar image patches.\")\n",
    "        print(\n",
    "            \"Clusters indicate that the model has learned meaningful representations!\"\n",
    "        )\n",
    "\n",
    "    except ImportError:\n",
    "        print(\"scikit-learn not available for t-SNE visualization.\")\n",
    "        print(\"Install with: pip install scikit-learn\")\n",
    "else:\n",
    "    print(\"Embeddings not available for visualization.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "Now that you have a pretrained self-supervised model, here are some ways you can use it:\n",
    "\n",
    "### 1. Fine-tuning for specific tasks\n",
    "```python\n",
    "# Load your pretrained model\n",
    "model = geoai.load_lightly_pretrained_model(\n",
    "    model_path=\"path/to/your/model.pt\",\n",
    "    model_architecture=\"torchvision/resnet50\"\n",
    ")\n",
    "\n",
    "# Replace the final layer for your specific task\n",
    "# For example, for binary classification:\n",
    "import torch.nn as nn\n",
    "model.fc = nn.Linear(model.fc.in_features, 2)  # 2 classes\n",
    "\n",
    "# Fine-tune with your labeled data using standard PyTorch training\n",
    "```\n",
    "\n",
    "### 2. Feature extraction\n",
    "```python\n",
    "# Use the model as a feature extractor\n",
    "model.eval()\n",
    "# Remove the final classification layer\n",
    "feature_extractor = nn.Sequential(*list(model.children())[:-1])\n",
    "\n",
    "# Extract features for any new images\n",
    "with torch.no_grad():\n",
    "    features = feature_extractor(your_image_tensor)\n",
    "```\n",
    "\n",
    "### 3. Similarity search\n",
    "```python\n",
    "# Use embeddings for finding similar images\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "\n",
    "# Find most similar images to a query image\n",
    "similarities = cosine_similarity([query_embedding], embeddings)\n",
    "most_similar_indices = similarities.argsort()[0][-5:]  # Top 5 similar\n",
    "```\n",
    "\n",
    "### 4. Different architectures and methods\n",
    "Try different combinations for your specific use case:\n",
    "\n",
    "**For CNN models (ResNet, EfficientNet):**\n",
    "```python\n",
    "# SimCLR method (recommended for CNNs)\n",
    "model_path = geoai.lightly_train_model(\n",
    "    data_dir=\"path/to/images\",\n",
    "    output_dir=\"output\",\n",
    "    model=\"torchvision/resnet50\",\n",
    "    method=\"simclr\",\n",
    "    epochs=100\n",
    ")\n",
    "\n",
    "# Or DINO method (also works with CNNs)\n",
    "model_path = geoai.lightly_train_model(\n",
    "    data_dir=\"path/to/images\",\n",
    "    output_dir=\"output\",\n",
    "    model=\"timm/efficientnet_b0\",\n",
    "    method=\"dino\",\n",
    "    epochs=100\n",
    ")\n",
    "```\n",
    "\n",
    "**For Vision Transformer (ViT) models:**\n",
    "```python\n",
    "# DINOv2 (recommended for ViTs, excellent for geospatial imagery)\n",
    "model_path = geoai.lightly_train_model(\n",
    "    data_dir=\"path/to/images\",\n",
    "    output_dir=\"output\",\n",
    "    model=\"timm/vit_base_patch16_224\",\n",
    "    method=\"dinov2\",\n",
    "    epochs=100\n",
    ")\n",
    "```\n",
    "\n",
    "**Model architectures:**\n",
    "- `\"torchvision/resnet50\"` - Good general purpose CNN\n",
    "- `\"torchvision/resnet101\"` - Larger CNN for more complex features\n",
    "- `\"timm/efficientnet_b0\"` - Efficient CNN architecture\n",
    "- `\"timm/vit_base_patch16_224\"` - Vision Transformer (use with dinov2)\n",
    "\n",
    "**Self-supervised methods:**\n",
    "- `\"simclr\"` - Contrastive learning, works with CNNs\n",
    "- `\"dino\"` - Works with both CNNs and ViTs\n",
    "- `\"dinov2\"` - Advanced method, requires ViT models\n",
    "- `\"dinov2_distillation\"` - Enhanced DINOv2, requires ViT models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "## Cleanup (optional)\n",
    "\n",
    "Remove temporary files if needed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to clean up temporary files\n",
    "# import shutil\n",
    "# shutil.rmtree(data_dir, ignore_errors=True)\n",
    "# shutil.rmtree(output_dir, ignore_errors=True)\n",
    "# shutil.rmtree(embeddings_dir, ignore_errors=True)\n",
    "# print(\"Temporary files cleaned up.\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "geoai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
